package cn.teamstack.logger.websocket;

import cn.teamstack.logger.config.SshSessionConfig;
import cn.teamstack.logger.enums.SshSessionVerifyEnum;
import cn.teamstack.logger.model.TailLogCmd;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.jcraft.jsch.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.ObjectUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import javax.annotation.Resource;
import javax.swing.*;
import java.io.InputStream;
import java.util.Properties;

/**
 * @author zhouli
 * @date 2017/9/18
 */
public class WebSocketHandler extends TextWebSocketHandler {
    private static final Logger log = LoggerFactory.getLogger(TextWebSocketHandler.class);
    @Resource
    private SshSessionConfig sshSessionConfig;

    @Override
    public void handleTextMessage(WebSocketSession webSocketSession, TextMessage textMessage) {
        String message = textMessage.getPayload();
        TailLogCmd tailLogCmd = JSON.parseObject(message, TailLogCmd.class);
        try {
            Session session = getSession(tailLogCmd);
            UserInfo ui = new AbstractMyUserInfo() {
                @Override
                public void showMessage(String message) {
                    JOptionPane.showMessageDialog(null, message);
                }

                @Override
                public boolean promptYesNo(String message) {
                    return true;
                }
            };
            session.setUserInfo(ui);
            // making a connection with timeout.
            session.connect(30000);
            Channel channel = session.openChannel("exec");
            //控制权限，防止误操作
            ((ChannelExec) channel).setCommand(String.format("tail -f %s",tailLogCmd.getContent()));
            channel.setInputStream(null);
            ((ChannelExec) channel).setErrStream(System.err);
            InputStream inputStream = channel.getInputStream();
            channel.connect();
            ProcessThread thread = new ProcessThread(inputStream, webSocketSession, session, channel);
            thread.start();

        } catch (Exception e) {
            log.error("JSch session error",e);
        }
    }

    /**
     * 获取WS Session
     * @return
     */
    private Session getSession(TailLogCmd tailLogCmd) throws JSchException {
        JSch jsch = new JSch();
        Session session = null;
        String sshVerify;
        String sshIP = null;
        String sshPort = null;
        String sshUsername = null;
        String sshPassword = null;
        String keysFilePath = null;
        String keysPassphrase = null;

        if (!ObjectUtils.isEmpty(tailLogCmd.sshInfo) && !StringUtils.isEmpty(tailLogCmd.sshInfo.ip)) {
            log.info("use ws agent config:{}", JSONObject.toJSON(tailLogCmd));
            //可以优先使用配置的ssh配置信息
            sshVerify = tailLogCmd.getSshInfo().getVerify();
            sshIP = tailLogCmd.getSshInfo().getIp();
            sshPort = tailLogCmd.getSshInfo().getPort();
            sshUsername = tailLogCmd.getSshInfo().getUsername();
            sshPassword = tailLogCmd.getSshInfo().getPassword();
            keysFilePath = tailLogCmd.getSshInfo().getKeysFilePath();
            keysPassphrase = tailLogCmd.getSshInfo().getKeysPassphrase();

        } else {
            //使用代理程序默认配置的信息
            log.info("use default config ,sshSessionConfig:{}",JSONObject.toJSON(sshSessionConfig));
            sshVerify = sshSessionConfig.getVerify();
            sshIP = sshSessionConfig.getIp();
            sshPort = sshSessionConfig.getPort();
            sshUsername = sshSessionConfig.getUsername();
            sshPassword = sshSessionConfig.getPassword();
            keysFilePath = sshSessionConfig.getKeysFilePath();
            keysPassphrase = sshSessionConfig.getKeysPassphrase();
        }
        SshSessionVerifyEnum sessionVerify = SshSessionVerifyEnum.getSSHSessionVerifyByName(sshVerify);
        if(SshSessionVerifyEnum.KEYS_FILE.equals(sessionVerify)){
            // 设置私钥文件路径及密码
            if (!StringUtils.isEmpty(keysPassphrase)) {
                jsch.addIdentity(keysFilePath, keysPassphrase);
            } else {
                jsch.addIdentity(keysFilePath);
            }
            log.info("WebSocketHandler getSession keyFilePath",keysFilePath);
        }
        session = jsch.getSession(sshUsername, sshIP, Integer.parseInt(sshPort));
        if(!StringUtils.isEmpty(sshPassword)){
            session.setPassword(sshPassword);
        }

        return session;
    }

    /**
     * 针对sftp的配置, 如果要sftp连接需要添加此项配置
     * @param session
     */
    private void configSession(Session session){
        Properties sshConfig = new Properties();
        sshConfig.put("StrictHostKeyChecking", "no");
        session.setConfig("kex", "diffie-hellman-group1-sha1");
        session.setConfig(sshConfig);
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        log.debug("Opened new session in instance " + this);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {

    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception)
            throws Exception {
        session.close(CloseStatus.SERVER_ERROR);
    }

    public static abstract class AbstractMyUserInfo
            implements UserInfo, UIKeyboardInteractive {
        @Override
        public String getPassword() {
            return null;
        }

        @Override
        public boolean promptYesNo(String str) {
            return false;
        }

        @Override
        public String getPassphrase() {
            return null;
        }

        @Override
        public boolean promptPassphrase(String message) {
            return false;
        }

        @Override
        public boolean promptPassword(String message) {
            return false;
        }

        @Override
        public void showMessage(String message) {
        }

        @Override
        public String[] promptKeyboardInteractive(String destination,
                                                  String name,
                                                  String instruction,
                                                  String[] prompt,
                                                  boolean[] echo) {
            return null;
        }
    }
}
